#include "videorecorder_mpp.h"

#include <QDebug>

VideoRecorder_MPP::VideoRecorder_MPP(QObject *parent) : BaseRecorder{parent} {}

VideoRecorder_MPP::~VideoRecorder_MPP() { uninit(); }

bool VideoRecorder_MPP::init() {
  initConf();
  // 缓存申请
  MPP_RET ret = mpp_buffer_group_get_internal(&buf_grp, MPP_BUFFER_TYPE_DRM);
  if (MPP_OK != ret) {
    qDebug() << "get buff group failed";
    return false;
  }

  ret = mpp_buffer_get(buf_grp, &frm_buf,
                       m_encoderConf.frame_size + m_encoderConf.header_size);
  if (MPP_OK != ret) {
    qDebug() << "get frame buff failed";
    return false;
  }

  ret = mpp_buffer_get(buf_grp, &pkt_buf, m_encoderConf.frame_size);
  if (MPP_OK != ret) {
    qDebug() << "get packet buff failed";
    return false;
  }

  // 创建一个mpp
  ret = mpp_create(&m_ctx, &m_pMpi);
  if (MPP_OK != ret) {
    return false;
  }

  // 设置为阻塞
  int timeout = MPP_TIMEOUT_BLOCK;
  ret = m_pMpi->control(m_ctx, MPP_SET_OUTPUT_TIMEOUT, &timeout);
  if (MPP_OK != ret) {
    return false;
  }

  // 初始化mpp
  ret = mpp_init(m_ctx, MppCtxType::MPP_CTX_ENC,
                 MppCodingType::MPP_VIDEO_CodingAVC);
  if (MPP_OK != ret) {
    return false;
  }

  // if (!initFrame()) {
  //   return false;
  // }

  // 设置配置
  if (!setupConf()) {
    return false;
  }

  return true;
}

void VideoRecorder_MPP::uninit() {
  if (m_frame) {
    mpp_frame_deinit(&m_frame);
    m_frame = nullptr;
  }
  if (m_ctx) {
    mpp_destroy(m_ctx);
    m_ctx = nullptr;
  }

  if (m_enc_cfg) {
    mpp_enc_cfg_deinit(m_enc_cfg);
    m_enc_cfg = nullptr;
  }

  m_outFile.close();
}

bool VideoRecorder_MPP::openFile(const QString &basename) {
  m_outFile.setFileName(basename);

  if (!m_outFile.open(QIODevice::WriteOnly | QIODevice::Truncate)) {
    return false;
  }

  // 写文件头
  MppPacket packet = NULL;
  mpp_packet_init_with_buffer(&packet, pkt_buf);
  mpp_packet_set_length(packet, 0);

  int ret = m_pMpi->control(m_ctx, MPP_ENC_GET_HDR_SYNC, packet);
  if (ret) {
    return false;
  } else {
    /*get and write sps/pps for H.264 */
    void *ptr = mpp_packet_get_pos(packet);
    size_t len = mpp_packet_get_length(packet);
    m_outFile.write(static_cast<char *>(ptr), len);
  }

  mpp_packet_deinit(&packet);
  return true;
}

void VideoRecorder_MPP::onImage(std::shared_ptr<QImage> img) {
  cv::Mat yuvFrame;
  image2Yuv420(img, &yuvFrame);
  writeToFile(yuvFrame);
}

// 设置配置参数的数值
void VideoRecorder_MPP::initConf() {
  // 帧大小
  m_encoderConf.frame_size = MPP_ALIGN(m_encoderConf.hor_stride, 1) *
                             MPP_ALIGN(m_encoderConf.ver_stride, 1) * 3 / 2;

  m_encoderConf.header_size = 0;
}

bool VideoRecorder_MPP::setupConf() {
  MPP_RET ret = MPP_OK;
  ret = mpp_enc_cfg_init(&m_enc_cfg);
  if (MPP_OK != ret) {
    qDebug() << "config init failed";
    return false;
  }

  // 镜像
  mpp_enc_cfg_set_s32(m_enc_cfg, "prep:mirroring",
                      2);  // 0 disable， 1 水平，2 垂直
  // 旋转
  mpp_enc_cfg_set_s32(m_enc_cfg, "prep:rotation",
                      MppEncRotationCfg_e::MPP_ENC_ROT_180);

  // config the ctx
  mpp_enc_cfg_set_s32(m_enc_cfg, "prep:width", m_encoderConf.frame_width);
  mpp_enc_cfg_set_s32(m_enc_cfg, "prep:height", m_encoderConf.frame_height);
  mpp_enc_cfg_set_s32(m_enc_cfg, "prep:hor_stride", m_encoderConf.hor_stride);
  mpp_enc_cfg_set_s32(m_enc_cfg, "prep:ver_stride", m_encoderConf.ver_stride);
  mpp_enc_cfg_set_s32(m_enc_cfg, "prep:format", m_encoderConf.format);

  mpp_enc_cfg_set_s32(m_enc_cfg, "rc:mode", m_encoderConf.rc_mode);

  // fix input/output frame rate
  mpp_enc_cfg_set_s32(m_enc_cfg, "rc:fps_in_flex", m_encoderConf.fps_in_flex);
  mpp_enc_cfg_set_s32(m_enc_cfg, "rc:fps_in_num", m_encoderConf.fps_in_num);
  mpp_enc_cfg_set_s32(m_enc_cfg, "rc:fps_in_denorm", m_encoderConf.fps_in_den);
  mpp_enc_cfg_set_s32(m_enc_cfg, "rc:fps_out_flex", m_encoderConf.fps_out_flex);
  mpp_enc_cfg_set_s32(m_enc_cfg, "rc:fps_out_num", m_encoderConf.fps_out_num);
  mpp_enc_cfg_set_s32(m_enc_cfg, "rc:fps_out_denorm",
                      m_encoderConf.fps_out_den);
  mpp_enc_cfg_set_s32(m_enc_cfg, "rc:gop",
                      m_encoderConf.gop_len
                          ? m_encoderConf.gop_len
                          : m_encoderConf.fps_out_num);  // 关键帧设置

  /*drop frame or not when bitrate overflow */
  mpp_enc_cfg_set_u32(m_enc_cfg, "rc:drop_mode", MPP_ENC_RC_DROP_FRM_DISABLED);
  mpp_enc_cfg_set_u32(m_enc_cfg, "rc:drop_thd", 20); /*20% of max bps */
  mpp_enc_cfg_set_u32(m_enc_cfg, "rc:drop_gap",
                      1); /*Do not continuous drop frame */

  /*setup bitrate for different rc_mode */
  mpp_enc_cfg_set_s32(m_enc_cfg, "rc:bps_target", m_encoderConf.bps);

  // vbr
  /*VBR mode has wide bound */
  mpp_enc_cfg_set_s32(m_enc_cfg, "rc:bps_max",
                      m_encoderConf.bps_max ? m_encoderConf.bps_max
                                            : m_encoderConf.bps * 17 / 16);
  mpp_enc_cfg_set_s32(m_enc_cfg, "rc:bps_min",
                      m_encoderConf.bps_min ? m_encoderConf.bps_min
                                            : m_encoderConf.bps * 10 / 16);

  // 配置qp
  mpp_enc_cfg_set_s32(m_enc_cfg, "rc:qp_init", 26);
  mpp_enc_cfg_set_s32(m_enc_cfg, "rc:qp_max", 51);
  mpp_enc_cfg_set_s32(m_enc_cfg, "rc:qp_min", 10);
  mpp_enc_cfg_set_s32(m_enc_cfg, "rc:qp_max_i", 51);
  mpp_enc_cfg_set_s32(m_enc_cfg, "rc:qp_min_i", 10);
  mpp_enc_cfg_set_s32(m_enc_cfg, "rc:qp_ip", 2);

  mpp_enc_cfg_set_s32(m_enc_cfg, "codec:type", m_encoderConf.codecype);

  /*
   *H.264 profile_idc parameter
   *66  - Baseline profile
   *77  - Main profile
   *100 - High profile
   */
  mpp_enc_cfg_set_s32(m_enc_cfg, "h264:profile", 100);

  /*
   *H.264 level_idc parameter
   *10 / 11 / 12 / 13    - qcif@15fps / cif@7.5fps / cif@15fps / cif@30fps
   *20 / 21 / 22         - cif@30fps / half-D1@@25fps / D1@12.5fps
   *30 / 31 / 32         - D1@25fps / 720p@30fps / 720p@60fps
   *40 / 41 / 42         - 1080p@30fps / 1080p@30fps / 1080p@60fps
   *50 / 51 / 52         - 4K@30fps
   */
  mpp_enc_cfg_set_s32(m_enc_cfg, "h264:level", 42);
  mpp_enc_cfg_set_s32(m_enc_cfg, "h264:cabac_en", 1);
  mpp_enc_cfg_set_s32(m_enc_cfg, "h264:cabac_idc", 0);
  mpp_enc_cfg_set_s32(m_enc_cfg, "h264:trans8x8", 0);

  ret = m_pMpi->control(m_ctx, MPP_ENC_SET_CFG, m_enc_cfg);
  if (ret) {
    qDebug() << "set config failed";
    return false;
  }

  ret = m_pMpi->control(m_ctx, MPP_ENC_SET_SEI_CFG, &sei_mode);
  if (ret) {
    qDebug() << "mpi control enc set sei cfg failed";
    return false;
  }

  ret = m_pMpi->control(m_ctx, MPP_ENC_SET_HEADER_MODE, &header_mode);
  if (ret) {
    qDebug() << "mpi control enc set header mode failed";
    return false;
  }
  return true;
}

// QImage转换成yuv
bool VideoRecorder_MPP::image2Yuv420(std::shared_ptr<QImage> img,
                                     cv::Mat *yuvFrame) {
  cv::Mat mat = cv::Mat(img->height(), img->width(), CV_8UC3,
                        (void *)img->constBits(), img->bytesPerLine());

  cv::cvtColor(mat, *yuvFrame, cv::COLOR_RGB2YUV_I420);

  return true;
}

bool VideoRecorder_MPP::writeToFile(const cv::Mat &yuvFrame) {
  int ret = mpp_frame_init(&m_frame);
  if (ret) {
    return false;
  }

  mpp_frame_set_width(m_frame, m_encoderConf.frame_width);
  mpp_frame_set_height(m_frame, m_encoderConf.frame_height);
  mpp_frame_set_hor_stride(m_frame, m_encoderConf.hor_stride);
  mpp_frame_set_ver_stride(m_frame, m_encoderConf.ver_stride);
  mpp_frame_set_fmt(m_frame, m_encoderConf.format);
  mpp_frame_set_eos(m_frame, 0);

  ret = mpp_buffer_write(frm_buf, 0, yuvFrame.data, m_encoderConf.frame_size);

  if (ret) {
    return false;
  }

  MppPacket packet = nullptr;
  mpp_frame_set_buffer(m_frame, frm_buf);
  MppMeta meta = mpp_frame_get_meta(m_frame);

  // packet
  mpp_packet_init_with_buffer(&packet, pkt_buf);
  mpp_packet_set_length(packet, 0);
  mpp_meta_set_packet(meta, KEY_OUTPUT_PACKET, packet);

  ret = m_pMpi->encode_put_frame(m_ctx, m_frame);
  mpp_frame_deinit(&m_frame);
  m_frame = nullptr;
  if (ret) {
    return false;
  }

  RK_U32 eoi = 1;

  // 接收
  do {
    ret = m_pMpi->encode_get_packet(m_ctx, &packet);
    if (ret) {
      return false;
    }

    if (packet) {
      void *ptr = mpp_packet_get_pos(packet);
      size_t len = mpp_packet_get_length(packet);
      m_outFile.write(static_cast<char *>(ptr), len);
      if (mpp_packet_is_partition(packet)) {
        eoi = mpp_packet_is_eoi(packet);
      }
    } else {
      //   printf("there is no packet\n");
    }

    mpp_packet_deinit(&packet);

  } while (!eoi);

  return true;
}

bool VideoRecorder_MPP::initFrame() {
  int ret = mpp_frame_init(&m_frame);
  if (ret) {
    return false;
  }

  mpp_frame_set_width(m_frame, m_encoderConf.frame_width);
  mpp_frame_set_height(m_frame, m_encoderConf.frame_height);
  mpp_frame_set_hor_stride(m_frame, m_encoderConf.hor_stride);
  mpp_frame_set_ver_stride(m_frame, m_encoderConf.ver_stride);
  mpp_frame_set_fmt(m_frame, m_encoderConf.format);
  mpp_frame_set_eos(m_frame, 0);

  return true;
}
